package io.adrainty.bolee.framework.gateway.decorator;

import io.netty.buffer.UnpooledByteBufAllocator;
import lombok.Getter;
import lombok.NonNull;
import lombok.extern.slf4j.Slf4j;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.core.io.buffer.NettyDataBufferFactory;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpRequestDecorator;
import reactor.core.publisher.Flux;
import reactor.core.scheduler.Schedulers;

/**
 * @author AdRainty
 * @version V1.0.0
 * @since 2025/7/27 下午10:41
 */

@Slf4j
@Getter
public class CacheServerHttpRequestDecorator extends ServerHttpRequestDecorator {

    private byte[] bytes;

    public CacheServerHttpRequestDecorator(ServerHttpRequest delegate) {
        super(delegate);
    }

    @NonNull
    @Override
    public Flux<DataBuffer> getBody() {
        // 获得父类获得请求体信息
        return super.getBody()
                // 切换单线程防止阻塞
                .publishOn(Schedulers.single())
                // 重写body可多次读取
                .map(this::cache);
    }

    private DataBuffer cache(DataBuffer dataBuffer) {
        // 先构建一个与dataBuffer长度相同的byte数组
        bytes = new byte[dataBuffer.readableByteCount()];
        // 从dataBuffer中读出数据存入当前bytes中
        dataBuffer.read(bytes);
        // 释放掉内存
        DataBufferUtils.release(dataBuffer);
        // 放回到当前的NettyDataBufferFactory，用于下游系统的使用
        NettyDataBufferFactory nettyDataBufferFactory = new NettyDataBufferFactory(new UnpooledByteBufAllocator(false));
        return nettyDataBufferFactory.wrap(bytes);
    }

}
